#!/usr/bin/env python3.9
# coding=utf-8

"""
Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
Description: python api test.
Author: MindX SDK
Create: 2022
History: NA
"""

import logging
import time
from multiprocessing import JoinableQueue, Queue
import traceback
import numpy as np
import cv2
from mindx.sdk import base
from mindx.sdk.base import Image, ImageProcessor, Size, Tensor

from utils import motor_post_process, common
from config import *
import logging.config
logging.config.fileConfig("logger.conf")
logger = logging.getLogger("pipeline")


class MotorRec(object):
    def __init__(self, rec_model_path, batch_size, device_id, pid):
        self.rec_model = base.model(rec_model_path, deviceId=device_id)
        self.batch_size = batch_size
        self.device_id = device_id
        self.pid = pid

        self.INFER_BREAK_WAIT_TIME = 5

    def rec_preprocess(self, img):
        height = 256
        width = 192
        # resize
        img = cv2.resize(img, (width, height))
        img = img[:, :, ::-1].transpose(2, 0, 1)  # HWC转CHW

        # 归一化
        img = img / 255

        # 标准化
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        mean = np.array(mean).reshape(3, 1, 1)
        std = np.array(std).reshape(3, 1, 1)
        img = (img - mean) / std

        img = img.astype(np.float32)
        img = np.expand_dims(img, axis=0)  # NCHW

        return img

    def infer(self, input: Queue):
        try:
            ImageProcessor(self.device_id)
            count = 0
            img_arr_list = []
            input_data_list = []
            while True:
                input_data = input.get()
                input_data_list.append(input_data)
                # img_ndarray = input_data["image"]
                img_ndarray = input_data["det"][0]
                # img_ndarray = input.get(block=True)
                if img_ndarray is None:
                    continue

                count += 1

                # 3. rec model preprocess
                img_ndarray = self.rec_preprocess(img_ndarray)  # HWC->NCHW
                img_arr_list.append(img_ndarray)

                if len(img_arr_list) < self.batch_size:
                    continue

                img_tensor = Tensor(np.squeeze(np.array(img_arr_list), axis=1))
                img_tensor.to_device(self.device_id)
                img_tensor_vec = [img_tensor]

                # 4. rec model infer
                output_tensors = self.rec_model.infer(img_tensor_vec)[0]
                output_tensors.to_host()
                output_tensors = np.array(output_tensors)

                # 5. rec model postprocess
                for i, output_tensor in enumerate(output_tensors):

                    feature = motor_post_process.rec_postprocess(output_tensors[i])
                    indexes = feature > 0.5
                    labels = np.array(MODEL_BODY_LABEL)[indexes]
                    probes = feature[indexes]
                    _rs = {}
                    for j, label in enumerate(labels):
                        attr = label.split('-')[0]
                        value = label.split('-')[1]
                        if attr not in _rs:
                            _rs[attr] = {}
                        _rs[attr][value] = float(probes[j])
                    whether_attrs = ['卷发-是',
                                     '戴口罩-是',
                                     '戴眼镜-是',
                                     '手推车-是',
                                     '抱小孩-是',
                                     '打伞-是',
                                     '手拉箱-是',
                                     '单肩包-是',
                                     '手提包-是',
                                     '双肩包-是',
                                     '斜挎包-是', ]
                    out_attrs = ['手拉箱', '单肩包', '手提包', '双肩包', '斜挎包']
                    other_attrs = ['性别', '年龄', '头发长度', '发型', '帽子款式', '帽子颜色',
                                   '上衣颜色', '上衣长度', '上衣纹理', '下身颜色', '下身长度', ]
                    for attr in whether_attrs:
                        index = MODEL_BODY_LABEL.index(attr)
                        attr = attr.split('-')[0]
                        if attr not in _rs:
                            _rs[attr] = {'否': 1 - feature[index]}  # 如没有戴帽子，则要返回没有戴帽子的概率
                    for out_attr in out_attrs:
                        if out_attr in _rs:
                            if '否' in _rs[out_attr] and out_attr + '颜色' in _rs:
                                _rs.pop(out_attr + '颜色')
                                out_attrs.remove(out_attr)
                            elif '是' in _rs[out_attr] and out_attr + '颜色' not in _rs:
                                _rs[out_attr + '颜色'] = {'无法判断': 0.99}
                    for other_attr in other_attrs:
                        if other_attr not in _rs:
                            _rs[other_attr] = {'无法判断': 0.99}
                    post_data = input_data_list[i]
                    post_data["attr"] = _rs
                    doc = self.gen_doc(post_data)
                    # 推送kafka
                    if doc is not None:
                        common.send_kafka(KAFKA_MOTOR_TOPIC, doc)

                img_arr_list = []
                input_data_list = []
        except Exception as e:
            print("process2_pattern_rec failed", repr(e))
            traceback.print_exc()

    def gen_doc(self, post_data):
        xb = list(post_data["attr"]["性别"].keys())[0]
        age = list(post_data["attr"]["年龄"].keys())[0]
        syys = list(post_data["attr"]["上衣颜色"].keys())[0]
        sycd = list(post_data["attr"]["上衣长度"].keys())[0]
        xsys = list(post_data["attr"]["下身颜色"].keys())[0]
        xscd = list(post_data["attr"]["下身长度"].keys())[0]
        mzks = list(post_data["attr"]["帽子款式"].keys())[0]
        dmz = "否"
        if mzks != "无法判断":
            dmz = "是"
        mzys = list(post_data["attr"]["帽子颜色"].keys())[0]
        dkz = list(post_data["attr"]["戴口罩"].keys())[0]
        dyj = list(post_data["attr"]["戴眼镜"].keys())[0]
        id = common.get_id()

        small_pic = f"/motor/{id}_crop.jpg"
        big_pic = f"/motor/{id}_orig.jpg"
        tmp_pic = f"/motor/{id}_tmp.jpg"
        orig = post_data["image"]
        img = orig[0]
        x1, y1, x2, y2 = post_data["det"][1]
        x1 = int(x1)
        y1 = int(y1)
        x2 = int(x2)
        y2 = int(y2)

        small_img = img[y1 - 50: y2+ 50, x1 - 50: x2 + 50].copy()
        # cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 255), 1)
        try:
            common.upload_pic(UPLOAD_PATH + small_pic, small_img)
            common.upload_pic(UPLOAD_PATH + big_pic, img)
        except:
            return None
        doc = {
            "id": id,
            "obj_id": post_data["track_id"],
            "big_id": id,
            "point_id": post_data["video_id"],
            "point_name": post_data["video_name"],
            "capture_time": post_data["capture_time"],
            "small_url": NGINX_URL + small_pic,
            "big_url": NGINX_URL + big_pic,
            "feature": "",
            "vehicle_type": "二轮",
            "manned": "否",
            "hat": dmz,
            "glasses": dkz,
            "mask": dyj,
            "gender": xb,
            "age": age,
            "upper_color": syys,
            "upper_length": sycd,
            "hat_color": mzys,
        }
        return doc


def motor_recognize(rec_model_path, batch_size, device_id, pid,input: Queue, model_loading: list):
    model = MotorRec(rec_model_path, batch_size, device_id, pid)
    model_loading.append(True)
    model.infer(input)

